from apk import ResNet_model
from PIL import Image
import torch
from torchvision import transforms, datasets
import json
import numpy as np
import pandas as pd
import joblib
import csv

# 四个模型：ResNet、随机森林、k近邻、SVM
def resnet2predict(image_path):

    image_processing = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    # losd image
    img = Image.open(image_path)
    img = image_processing(img)
    # expand batch dimension
    img = torch.unsqueeze(img, dim=0)

    params_dict = torch.load("C:\\file_data\涉诈APP\\model_parameter\\resnet152_best.pth")
    net = ResNet_model.resnet152(num_classes=2)
    net.load_state_dict(params_dict)
    net.eval()
    try:
        json_file = open('C:\\file_data\涉诈APP\\model_parameter\\class_indices.json', 'r')
        class_indict = json.load(json_file)
    except Exception as e:
        print(e)
        exit(-1)

    with torch.no_grad():
        # predict class
        output = torch.squeeze(net(img))
        predict = torch.softmax(output, dim=0)
        predict_cla = torch.argmax(predict).numpy()
    print(class_indict[str(predict_cla)], predict[predict_cla].numpy())
    return class_indict[str(predict_cla)], predict[predict_cla].numpy()

def svm2predict(data):
    cla =["black", "gamble", "scam", 'sex', 'white']

    model = joblib.load("C:\\file_data\涉诈APP\\model_parameter\\SVM.pkl")

    # predict
    pre = model.predict_proba(data)
    res_pro = np.max(pre)
    res_idx = np.argmax(pre)
    if res_idx == 4:
        return [cla[res_idx], res_pro]
    else:
        return ["dengerous", cla[res_idx], res_pro]

def knn2predict(data):
    cla =["black", "gamble", "scam", 'sex', 'white']
    model = joblib.load("C:\\file_data\\涉诈APP\\model_parameter\\KNN.pkl")
    # predict
    pre = model.predict_proba(data)
    res_pro = np.max(pre)
    res_idx = np.argmax(pre)
    if res_idx == 4:
        return [cla[res_idx], res_pro]
    else:
        return ["dengerous", cla[res_idx], res_pro]

def rf2predict(data):
    cla =["black", "gamble", "scam", 'sex', 'white']

    model = joblib.load("C:\\file_data\\涉诈APP\\model_parameter\\RF.pkl")

    # predict
    pre = model.predict_proba(data)
    res_pro = np.max(pre)
    res_idx = np.argmax(pre)
    if res_idx == 4:
        return [cla[res_idx], res_pro]
    else:
        return ["dengerous", cla[res_idx], res_pro]


# 导入数据
def get_input_permissions(input_path):
    data = []
    with open(input_path, 'r') as f:
         _ = csv.reader(f)
         for row in _:
             if row:
                data.append(row)
             else:
                continue
    return data

'''
df = pd.read_csv("D:\\APP_Detection\\permission_data\\train_allpermissions2（131）.csv")
data = df.iloc[5, 1:]
res_data = []
for _ in data:
    res_data.append(_)
result = []
result.append(res_data)

svm_pre = svm2predict(result)
print(svm_pre)
knn_pre = knn2predict(result)
print(knn_pre)
rf_pre = rf2predict(result)
print(rf_pre)
'''